import math  # 导入标准库的math模块
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D


def radial_polynomial(n, m, rho):
    """计算径向多项式Z(n, m)的径向部分

    参数:
        n: Zernike多项式阶数，非负整数，必须满足n >= |m|
        m: Zernike模式编号，整数，|m| <= n 且n - |m|必须是偶数
        rho: 归一化径向坐标，范围[0, 1]

    返回:
        R(n, m)(rho)的值
    """
    if (n - abs(m)) % 2 != 0:
        return np.zeros_like(rho)

    R = 0
    for k in range(int((n - abs(m)) // 2) + 1):
        numerator = (-1) ** k * math.factorial(n - k)  # 使用math.factorial
        denominator = (
                math.factorial(k) *
                math.factorial((n + abs(m)) // 2 - k) *
                math.factorial((n - abs(m)) // 2 - k)
        )
        term = numerator / denominator * rho ** (n - 2 * k)
        R += term

    return R


def zernike_polynomial(n, m, rho, theta):
    """计算Zernike多项式Z(n, m)

    参数:
        n: Zernike多项式阶数，非负整数，必须满足n >= |m|
        m: Zernike模式编号，整数，|m| <= n 且n - |m|必须是偶数
        rho: 归一化径向坐标，范围[0, 1]
        theta: 角坐标，范围[0, 2π]

    返回:
        Z(n, m)(rho, theta)的值
    """
    R = radial_polynomial(n, m, rho)

    if m > 0:
        return np.sqrt(2) * R * np.cos(m * theta)
    elif m < 0:
        return np.sqrt(2) * R * np.sin(abs(m) * theta)
    else:  # m == 0
        return R


def generate_zernike_surface(n, m, resolution=100):
    """生成指定阶数和模式的Zernike多项式曲面

    参数:
        n: Zernike多项式阶数
        m: Zernike模式编号
        resolution: 网格分辨率

    返回:
        rho, theta: 极坐标网格
        Z: Zernike多项式值网格
    """
    rho = np.linspace(0, 1, resolution)
    theta = np.linspace(0, 2 * np.pi, resolution)
    rho_grid, theta_grid = np.meshgrid(rho, theta)

    Z = np.zeros(rho_grid.shape)
    for i in range(resolution):
        for j in range(resolution):
            Z[i, j] = zernike_polynomial(n, m, rho_grid[i, j], theta_grid[i, j])

    return rho_grid, theta_grid, Z


def visualize_zernike(n, m, resolution=100):
    """可视化指定阶数和模式的Zernike多项式

    参数:
        n: Zernike多项式阶数
        m: Zernike模式编号
        resolution: 网格分辨率
    """
    rho, theta, Z = generate_zernike_surface(n, m, resolution)
    x = rho * np.cos(theta)
    y = rho * np.sin(theta)

    fig = plt.figure(figsize=(14, 6))

    # 3D表面图
    ax1 = fig.add_subplot(121, projection='3d')
    surf = ax1.plot_surface(x, y, Z, cmap='coolwarm', linewidth=0, antialiased=False)
    ax1.set_title(f'Zernike Polynomial Z({n},{m}) (3D View)')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    ax1.set_zlabel('Z')
    fig.colorbar(surf, shrink=0.5, aspect=5)

    # 2D热图
    ax2 = fig.add_subplot(122)
    cax = ax2.contourf(x, y, Z, cmap='coolwarm', levels=50)
    ax2.set_title(f'Zernike Polynomial Z({n},{m}) (2D Heatmap)')
    ax2.set_xlabel('X')
    ax2.set_ylabel('Y')
    fig.colorbar(cax, ax=ax2, shrink=0.8, aspect=5)

    plt.tight_layout()
    plt.show()


# 生成并可视化不同阶的Zernike多项式
if __name__ == "__main__":
    # 可视化Z(1,1) - 倾斜模式
    # visualize_zernike(1, 1)

    # 可视化Z(3,1) - 此模式
    visualize_zernike(3, 1)

    # 可视化Z(4,0) - 四叶形模式
    # visualize_zernike(4, 0)